#
# Copyright (C) [2024] Xingyun Integrated Circuit, Inc.
#
#   GreenCode was a private technology asset of Xingyun Integrated Circuit， Inc （Confidential）
#   Author :  Shawn.Tan
#   Date : 2025.10.28
#
#   History : Initial Version 2025.10.28
#
#
import itertools
import os

from base.exception_handlers.ExceptionHandlerAssignment import (
    ExceptionHandlerAssignmentRequest,
)
from base.exception_handlers.ExceptionHandlerAssignmentParser import (
    ExceptionHandlerAssignmentParser,
)
from base.exception_handlers.ExceptionHandlerManager import (
    ExceptionHandlerManager,
)
from base.exception_handlers.MemoryBankHandlerRegistry import (
    MemoryBankHandlerRegistry,
)
from gpgpu.MemoryBank import MemoryBankGPGPU
from gpgpu.PrivilegeLevel import PrivilegeLevelGPGPU
from gpgpu.SecurityState import SecurityStateGPGPU
from gpgpu.exception_handlers.ExceptionClass import ExceptionClassGPGPU
from gpgpu.exception_handlers.HandlerSubroutineGenerator import (
    HandlerSubroutineGeneratorGPGPU,
)


class ExceptionHandlerManagerGPGPU(ExceptionHandlerManager):
    def __init__(self, gen_thread, factory):
        super().__init__(gen_thread, factory)

        self.mGenThread = gen_thread

        for mem_bank in MemoryBankGPGPU:
            self.memBankHandlerRegistryRepo.addMemoryBankHandlerRegistry(
                MemoryBankHandlerRegistry(gen_thread, mem_bank)
            )

    # return True if a particular exception can be redirected at a given
    # privilege level, based on General choices...
    def useTrapHandler(self, exceptionClassName, subexceptionClassName, privilegeLevel):
        if privilegeLevel == PrivilegeLevelGPGPU.M:
            # trap redirection supported (for now) only from machine-mode...
            pass
        else:
            return False

        # translate exception class name into redirect-trap choices tree...
        (class_prefix, exception_name) = exceptionClassName.split(".")
        exc_choice_trees = {
            "INSTRUCTION_ADDRESS_MISALIGNED": "Redirect Trap - Instruction address misaligned",
            "INSTRUCTION_ACCESS_FAULT": "Redirect Trap - Instruction access fault",
            "ILLEGAL_INSTRUCTION": "Redirect Trap - Illegal instruction",
            "BREAKPOINT": "Redirect Trap - Breakpoint",
            "LOAD_ADDRESS_MISALIGNED": "Redirect Trap - Load address misaligned",
            "LOAD_ACCESS_FAULT": "Redirect Trap - Load access fault",
            "STORE_AMO_ADDRESS_MISALIGNED": "Redirect Trap - Store/AMO address misaligned",
            "STORE_AMO_ACCESS_FAULT": "Redirect Trap - Store/AMO access fault",
            "ENV_CALL_FROM_U_MODE": "Redirect Trap - Environment call from U-mode",
            "ENV_CALL_FROM_S_MODE": "Redirect Trap - Environment call from S-mode",
            "INSTRUCTION_PAGE_FAULT": "Redirect Trap - Instruction page fault",
            "LOAD_PAGE_FAULT": "Redirect Trap - Load page fault",
            "STORE_AMO_PAGE_FAULT": "Redirect Trap - Store/AMO page fault",
        }

        use_trap_handler = False

        try:
            # make the choice...
            choices = self.mGenThread.getChoicesTreeInfo(
                exc_choice_trees[exception_name], "GeneralChoices"
            )
            use_trap_handler = self.pickWeighted(choices) == "DoRedirect"
        except KeyError:
            # not all exceptions can be redirected...
            pass

        return use_trap_handler

    def registerDefaultExceptionHandlers(self):
        assignment_file_path = self.getDefaultAssignmentFilePath(self.default_set_name)
        assignment_parser = ExceptionHandlerAssignmentParser()
        handler_assignments = assignment_parser.parseHandlerAssignments(assignment_file_path)

        trap_handler_module_name = None
        trap_handler_class_name = None

        for handler_assignment in handler_assignments.items():
            (
                (exception_class_name, _),
                (handler_module_name, handler_class_name),
            ) = handler_assignment
            if "TRAP_REDIRECTION" in exception_class_name:
                trap_handler_module_name = handler_module_name
                trap_handler_class_name = handler_class_name
                continue

            self._assignHandler(
                handler_assignment,
                trap_handler_module_name,
                trap_handler_class_name,
            )

        self._assignSpecialHandlers()

    def configureHandlerMemory(self):
        (exc_memory, mem_size) = self.queryHandlerSetMemory("0")

        self.debugPrint("[ExceptionHandlerManagerGPGPU] EXC MEMORY START: 0x%x" % exc_memory)
        repo = self.memBankHandlerRegistryRepo
        registry = repo.getMemoryBankHandlerRegistry(MemoryBankGPGPU.DEFAULT)
        registry.mStartAddr = exc_memory

        exception_bounds_info_set = {
            "Function": "UpdateExceptionBounds",
            "memory_bounds": "%s,%s"
            % (
                exc_memory,
                mem_size,
            ),
        }
        self.exceptionRequest("UpdateHandlerInfo", exception_bounds_info_set)

    def getDefaultAssignmentFilePath(self, defaultSetName):
        assignment_file_name = "default_%s_exception_handlers.json" % defaultSetName.lower()
        assignment_file_path = os.path.join(os.path.dirname(__file__), assignment_file_name)
        return assignment_file_path

    # initialize vector base address registers for current thread from
    # exception set addresses
    def initializeVectorBaseAddressRegisters(self):
        mtvec_val = self.thread_handler_set.getVectorBaseAddress(
            PrivilegeLevelGPGPU.M, SecurityStateGPGPU.DEFAULT
        )
        self.initializeRegister("mtvec", mtvec_val)

        stvec_val = self.thread_handler_set.getVectorBaseAddress(
            PrivilegeLevelGPGPU.S, SecurityStateGPGPU.DEFAULT
        )
        self.initializeRegister("stvec", stvec_val)

    def getExceptionClass(self, exceptionClassName):
        (
            exception_code_class_name,
            _,
            err_code_name,
        ) = exceptionClassName.partition(".")
        return ExceptionClassGPGPU[err_code_name]

    def getMemoryBanks(self):
        return tuple(MemoryBankGPGPU)

    def getPrivilegeLevelSecurityStateCombinations(self):
        combinations = itertools.product(PrivilegeLevelGPGPU, SecurityStateGPGPU)

        def filter_valid(aPrivLevelSecurityStateCombo):
            (priv_level, security_state) = aPrivLevelSecurityStateCombo
            valid = True
            if priv_level == PrivilegeLevelGPGPU.U:
                valid = False

            return valid

        return filter(filter_valid, combinations)

    def _assignHandler(self, handlerAssignment, trapHandlerModuleName, trapHandlerClassName):
        (
            (exception_class_name, subexception_class_name),
            (handler_module_name, handler_class_name),
        ) = handlerAssignment
        exception_class = self.getExceptionClass(exception_class_name)

        subexception_class = None
        if subexception_class_name is not None:
            subexception_class = self.getExceptionClass(subexception_class_name)

        for (
            priv_level,
            security_state,
        ) in self.getPrivilegeLevelSecurityStateCombinations():
            if (trapHandlerClassName is not None) and self.useTrapHandler(
                exception_class_name, subexception_class_name, priv_level
            ):
                handler_module_name = trapHandlerModuleName
                handler_class_name = trapHandlerClassName

            handler_assignment_request = ExceptionHandlerAssignmentRequest(
                exception_class,
                (priv_level,),
                (security_state,),
                handler_module_name,
                handler_class_name,
                aMemBank=None,
                aSubexcClass=subexception_class,
            )

            self.thread_handler_set.assignSynchronousExceptionHandler(handler_assignment_request)

    def _assignSpecialHandlers(self):
        for registry in self.memBankHandlerRegistryRepo.getMemoryBankHandlerRegistries():
            generator = HandlerSubroutineGeneratorGPGPU(
                self.genThread, self.factory, self.exceptions_stack
            )
            registry.mHandlerSubroutineGenerator = generator

            if self.fastMode():
                registry.registerExceptHandlerWithClassName(
                    "gpgpu.exception_handlers.FastExceptionHandlers",
                    "FastEmptyHandlerGPGPU",
                    self.factory,
                    self.exceptions_stack,
                )
                self.thread_handler_set.assignAsynchronousExceptionHandler("FastEmptyHandlerGPGPU")
            else:
                registry.registerExceptHandlerWithClassName(
                    "gpgpu.exception_handlers.AsynchronousHandlers",
                    "AsynchronousHandlerGPGPU",
                    self.factory,
                    self.exceptions_stack,
                )
                self.thread_handler_set.assignAsynchronousExceptionHandler(
                    "AsynchronousHandlerGPGPU"
                )
